import os

import pandas as pd


file_config = {
    "市场指数-国内指数日行情文件.csv": {
        "output": "index_daily.pkl",
        "converters": {"交易所指数代码": str},
        "rename": {
            "交易所指数代码": "index_id",
            "交易日期": "date",
            "指数回报率": "factor_return",
        },
        "keeping": ["date", "index_id", "factor_return"],
        "start": "2010-01-01",
    },
    "市场指数-指数基本信息文件.csv": {
        "output": "index_info.pkl",
        "converters": None,
    },
}

cur_dir = os.path.dirname(__file__)
all_files = os.listdir(cur_dir)
for file in all_files:
    if file.endswith(".csv"):
        data = pd.read_csv(
            os.path.join(cur_dir, file), converters=file_config[file]["converters"]
        )
        if "rename" in file_config[file].keys():
            data.rename(columns=file_config[file]["rename"], inplace=True)
        if "keeping" in file_config[file].keys():
            data = data[file_config[file]["keeping"]]
        if "start" in file_config[file].keys():
            data = data[data["date"] >= file_config[file]["start"]]
        data.to_pickle(os.path.join(cur_dir, file_config[file]["output"]), "gzip")
